from unittest import TestCase, main
from sklearn.datasets import load_iris
import pandas as pd
from sklearn.model_selection import train_test_split
import mglearn
from matplotlib import pyplot as plot
from sklearn.neighbors import KNeighborsClassifier as KNC
import numpy as np


class IrisTest(TestCase):
    def __init__(self, *args, **kwargs):
        super(IrisTest, self).__init__(*args, **kwargs)
        self.iris = load_iris()
        self.split_iris = train_test_split(self.iris.data, self.iris.target, random_state=0)

    def setUp(self) -> None:
        pass

    def test_load(self):
        for method_name in dir(self.iris):
            if method_name not in ('DESCR', 'data', 'target'):
                print('%s: %s' % (method_name, self.iris.__getattr__(method_name)))

    def test_show_plot(self):
        iris = self.iris
        xtr, xte, ytr, yte = self.split_iris
        print(f'xtr shape: {xtr.shape}, xte shape: {xte.shape}')
        print(f'ytr.shape: {ytr.shape}, yte shape: {yte.shape}')

        df = pd.DataFrame(xtr, columns=iris.feature_names)
        pd.plotting.scatter_matrix(df, c=ytr, figsize=(15, 15), marker='o',
                                   hist_kwds={'bins': 20}, s=60, alpha=.8, cmap=mglearn.cm3)
        plot.show()

    def test_knn(self):
        # prepare data
        xtr, xte, ytr, yte = self.split_iris
        knn = KNC(n_neighbors=1)
        knn.fit(xtr, ytr)
        print(f'np.mean predict: {np.mean(yte == knn.predict(xte))}')
        print(f'knn.score predict: {knn.score(xte, yte)}')


if __name__ == '__main__':
    main()
